import cv2
import mediapipe as mp
import numpy as np

# 初始化 MediaPipe Pose
mp_pose = mp.solutions.pose
mp_drawing = mp.solutions.drawing_utils

# 定义火柴人连接点
stickman_connections = [
    (mp_pose.PoseLandmark.NOSE, mp_pose.PoseLandmark.LEFT_SHOULDER),
    (mp_pose.PoseLandmark.NOSE, mp_pose.PoseLandmark.RIGHT_SHOULDER),
    (mp_pose.PoseLandmark.LEFT_SHOULDER, mp_pose.PoseLandmark.RIGHT_SHOULDER),
    (mp_pose.PoseLandmark.LEFT_SHOULDER, mp_pose.PoseLandmark.LEFT_ELBOW),
    (mp_pose.PoseLandmark.RIGHT_SHOULDER, mp_pose.PoseLandmark.RIGHT_ELBOW),
    (mp_pose.PoseLandmark.LEFT_ELBOW, mp_pose.PoseLandmark.LEFT_WRIST),
    (mp_pose.PoseLandmark.RIGHT_ELBOW, mp_pose.PoseLandmark.RIGHT_WRIST),
    (mp_pose.PoseLandmark.LEFT_SHOULDER, mp_pose.PoseLandmark.LEFT_HIP),
    (mp_pose.PoseLandmark.RIGHT_SHOULDER, mp_pose.PoseLandmark.RIGHT_HIP),
    (mp_pose.PoseLandmark.LEFT_HIP, mp_pose.PoseLandmark.RIGHT_HIP),
    (mp_pose.PoseLandmark.LEFT_HIP, mp_pose.PoseLandmark.LEFT_KNEE),
    (mp_pose.PoseLandmark.RIGHT_HIP, mp_pose.PoseLandmark.RIGHT_KNEE),
    (mp_pose.PoseLandmark.LEFT_KNEE, mp_pose.PoseLandmark.LEFT_ANKLE),
    (mp_pose.PoseLandmark.RIGHT_KNEE, mp_pose.PoseLandmark.RIGHT_ANKLE),
]

def draw_stickman(image, landmarks, image_width, image_height):
    """绘制火柴人图像"""
    stickman_image = np.ones((image_height, image_width, 3), dtype=np.uint8) * 255  # 创建白色背景
    
    # 绘制连接线
    for connection in stickman_connections:
        start_idx = connection[0]
        end_idx = connection[1]
        
        if (landmarks.landmark[start_idx].visibility > 0.2 and 
            landmarks.landmark[end_idx].visibility > 0.2):
            start_point = (int(landmarks.landmark[start_idx].x * image_width),
                          int(landmarks.landmark[start_idx].y * image_height))
            end_point = (int(landmarks.landmark[end_idx].x * image_width),
                        int(landmarks.landmark[end_idx].y * image_height))
            
            cv2.line(stickman_image, start_point, end_point, (0, 0, 0), 2)
    
    # 绘制关节点
    for idx, landmark in enumerate(landmarks.landmark):
        if landmark.visibility > 0.2:
            x = int(landmark.x * image_width)
            y = int(landmark.y * image_height)
            cv2.circle(stickman_image, (x, y), 5, (0, 0, 255), -1)
    
    return stickman_image

def main():
    # 打开摄像头
    cap = cv2.VideoCapture(0)
    
    # 检查摄像头是否成功打开
    if not cap.isOpened():
        print("无法打开摄像头")
        exit()
    
    # 获取视频的宽度和高度
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    # 初始化 Pose 模型
    with mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5) as pose:
        while True:
            # 读取一帧视频
            ret, frame = cap.read()
            
            # 检查是否成功读取帧
            if not ret:
                print("无法获取帧")
                break
            
            # 转换为 RGB 格式
            image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            
            # 处理图像，检测姿势
            results = pose.process(image_rgb)
            
            # 在原始图像上绘制关节点和连接线
            annotated_image = frame.copy()
            if results.pose_landmarks:
                mp_drawing.draw_landmarks(
                    annotated_image,
                    results.pose_landmarks,
                    mp_pose.POSE_CONNECTIONS,
                    mp_drawing.DrawingSpec(color=(0, 255, 0), thickness=2, circle_radius=2),
                    mp_drawing.DrawingSpec(color=(0, 0, 255), thickness=2)
                )
                
                # 绘制火柴人
                stickman_image = draw_stickman(annotated_image, results.pose_landmarks, width, height)
                
                # 显示火柴人窗口
                cv2.imshow('Stickman', stickman_image)
            
            # 显示原始图像窗口
            cv2.imshow('Pose Tracking', annotated_image)
            
            # 按 'q' 键退出循环
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
    
    # 释放资源
    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()